#include <memory>
#include <string>
#include <vector>

#include <windows.h>
#include <BaseTsd.h>
#include <initguid.h>
#include <setupapi.h>
#include <winioctl.h>

#include "log.h"
#include "thirdparty/vioserial/public.h"
#include "VioSAPI.h"

const int LOOKUP_PORT_MAX = 32;
const int MAX_NAME_LEN = 512;

#define RECONNECT_RETRY_DELAY 100

VioSAPI::~VioSAPI()
{
    if (m_handle != INVALID_HANDLE_VALUE) {
        CloseHandle(m_handle);
        m_handle = INVALID_HANDLE_VALUE;
    }
}

BOOL VioSAPI::InitDev(INT hostPortIdx)
{
    if (m_bInitialized) {
        LOGW("VioSAPI already initialised!");
        return TRUE;
    }

    if (hostPortIdx < 0 || hostPortIdx >= LOOKUP_PORT_MAX) {
        LOGE("hostPortIdx out of range: {}", hostPortIdx);
        return FALSE;
    }

    if (!GetPortHandle(TRUE, hostPortIdx)) {
        LOGE("Failed to open vioserial port: {}", hostPortIdx);
        return FALSE;
    }

    if (m_handle == INVALID_HANDLE_VALUE) {
        LOGE("Cannot get port handle port: {}", hostPortIdx);
        return FALSE;
    }

    m_hostPortIdx = hostPortIdx;
    m_bInitialized = TRUE;
    return TRUE;
}

BOOL VioSAPI::PurgeIO()
{
    if (!CancelIoEx(m_handle, nullptr)) {
        DWORD dwErr = GetLastError();
        if (dwErr != ERROR_NOT_FOUND) {
            LOGE("Failed to cancel vioserial port: {}, CancelIoEx error: {}", m_hostPortIdx, dwErr);
            return FALSE;
        }
    }

    if (!PurgeComm(m_handle, PURGE_RXABORT | PURGE_RXCLEAR | PURGE_TXABORT | PURGE_TXCLEAR)) {
        LOGE("Failed to purge vioserial port: {}, PurgeComm error: {}", m_hostPortIdx, GetLastError());
        return FALSE;
    }

    return TRUE;
}

HANDLE VioSAPI::GetHandle()
{
    return m_handle;
}

BOOL VioSAPI::GetPortHandle(BOOL ovrl, INT hostPortIdx)
{
    HDEVINFO devInfo =
        SetupDiGetClassDevs((LPGUID)&GUID_VIOSERIAL_PORT, nullptr, nullptr, (DIGCF_PRESENT | DIGCF_DEVICEINTERFACE));
    if (devInfo == INVALID_HANDLE_VALUE) {
        LOGE("Cannot get class devices.");
        return FALSE;
    }

    for (INT idx = 0; idx < LOOKUP_PORT_MAX; ++idx) {
        SP_DEVICE_INTERFACE_DATA devData = { 0 };
        devData.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA);
        if (!SetupDiEnumDeviceInterfaces(devInfo, 0, (LPGUID)&GUID_VIOSERIAL_PORT, idx, &devData)) {
            continue;
        }

        ULONG reqLen = 0;
        SetupDiGetDeviceInterfaceDetail(devInfo, &devData, nullptr, 0, &reqLen, nullptr);
        std::vector<uint8_t> devDataBuf(reqLen);
        PSP_DEVICE_INTERFACE_DETAIL_DATA devIDData = (PSP_DEVICE_INTERFACE_DETAIL_DATA)devDataBuf.data();

        devIDData->cbSize = sizeof(SP_DEVICE_INTERFACE_DETAIL_DATA);
        ULONG length = reqLen;
        if (!SetupDiGetDeviceInterfaceDetail(devInfo, &devData, devIDData, length, &reqLen, nullptr)) {
            LOGE("Cannot get device interface details.");
            continue;
        }

        m_handle = CreateFile(devIDData->DevicePath, GENERIC_WRITE | GENERIC_READ, 0, nullptr,
            OPEN_EXISTING, ovrl ? FILE_FLAG_OVERLAPPED : FILE_ATTRIBUTE_NORMAL, nullptr);
        if (m_handle == INVALID_HANDLE_VALUE) {
            LOGD("Failed to CreateFile with error: {}", GetLastError());
            continue;
        }

        ULONG retLen = 0;
        BYTE buf[MAX_NAME_LEN + sizeof(VIRTIO_PORT_INFO)];
        if (!DeviceIoControl(m_handle, IOCTL_GET_INFORMATION, nullptr, 0, buf, sizeof(buf), &retLen, nullptr)) {
            LOGE("Virtio serial IOCTL_GET_INFORMATION failed {}", GetLastError());
            CloseHandle(m_handle);
            m_handle = INVALID_HANDLE_VALUE;
            continue;
        }

        PVIRTIO_PORT_INFO inf = (PVIRTIO_PORT_INFO)buf;
        if (inf->Id != hostPortIdx) {
            CloseHandle(m_handle);
            m_handle = INVALID_HANDLE_VALUE;
        } else {
            LOGI("Found virtio serial device with host port idx: {}", hostPortIdx);
            break;
        }
    }

    if (!SetupDiDestroyDeviceInfoList(devInfo)) {
        LOGE("SetupDiDestroyDeviceInfoList failed: {}", GetLastError());
    }
    return TRUE;
}

BOOL VioSAPI::Write(PVOID buf, DWORD &writeSz)
{
    OVERLAPPED ovl = { 0 };
    DWORD nbytes = 0;
    BOOL res = FALSE;

    ovl.hEvent = CreateEvent(nullptr, TRUE, FALSE, nullptr);
    if (!ovl.hEvent) {
        LOGE("Failed to create event for write to vioserial port: {}", m_hostPortIdx);
        return false;
    }

    res = WriteFile(m_handle, buf, writeSz, &nbytes, &ovl);
    if (!res) {
        if (GetLastError() != ERROR_IO_PENDING) {
            LOGE("Failed to write to vioserial port: {}", m_hostPortIdx);
        } else {
            if (GetOverlappedResult(m_handle, &ovl, &nbytes, TRUE)) {
                writeSz = nbytes;
                res = TRUE;
            } else {
                LOGE("Failed to GetOverlappedResult with error: {}", GetLastError());
            }
        }
    } else {
        writeSz = nbytes;
    }
    CloseHandle(ovl.hEvent);

    return res;
}

VioSOpStatus VioSAPI::ReadTimed(PVOID buf, DWORD &readSz, INT timeoutMs)
{
    OVERLAPPED ovl = {0};
    DWORD nbytes;
    BOOL res;
    VioSOpStatus ret = VioSOpStatus::SUCCESS;

    ovl.hEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
    if (!ovl.hEvent) {
        LOGE("Failed to create event for timed read of vioserial port: {}", m_hostPortIdx);
        return VioSOpStatus::FAILURE;
    }

    res = ReadFile(m_handle, buf, readSz, &nbytes, &ovl);
    if (!res) {
        if (GetLastError() != ERROR_IO_PENDING) {
            LOGE("Failed to timed read from vioserial port: {}", m_hostPortIdx);
            ret = VioSOpStatus::FAILURE;
            CloseHandle(ovl.hEvent);
            return ret;
        }

        if (!GetOverlappedResultEx(m_handle, &ovl, &nbytes, timeoutMs, TRUE)) {
            DWORD err = GetLastError();
            if (err == WAIT_TIMEOUT) {
                CancelIo(m_handle);
                ret = VioSOpStatus::TIMEOUT;
            } else {
                LOGE("Failed of timed read from vioserial port: {}, error: {}", m_hostPortIdx, err);
                ret = VioSOpStatus::FAILURE;
            }
        }
    }
    if (ret == VioSOpStatus::SUCCESS) {
        readSz = nbytes;
    }
    CloseHandle(ovl.hEvent);

    return ret;
}

VioSOpStatus VioSAPI::ReadAllTimed(PVOID buf, DWORD readSz, INT timeoutMs)
{
    DWORD nbytes;
    DWORD total = 0;
    VioSOpStatus ret;
    UINT8 *byteBuf = (UINT8 *)buf;

    do {
        nbytes = readSz - total;
        ret = ReadTimed((PVOID)byteBuf, nbytes, timeoutMs);
        if (ret != VioSOpStatus::SUCCESS) {
            break;
        }
        total += nbytes;
        byteBuf += nbytes;
    } while (total < readSz);
    return ret;
}

VioSOpStatus VioSAPI::WriteTimed(PVOID buf, DWORD &writeSz, INT timeoutMs)
{
    OVERLAPPED ovl = {0};
    DWORD nbytes;
    BOOL res;
    VioSOpStatus ret = VioSOpStatus::SUCCESS;

    ovl.hEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
    if (!ovl.hEvent) {
        LOGE("Failed to create event for write to vioserial port: {}", m_hostPortIdx);
        return VioSOpStatus::FAILURE;
    }

    res = WriteFile(m_handle, buf, writeSz, &nbytes, &ovl);
    if (!res) {
        if (GetLastError() != ERROR_IO_PENDING) {
            LOGE("Failed to timed write to vioserial port: {}", m_hostPortIdx);
            ret = VioSOpStatus::FAILURE;
            CloseHandle(ovl.hEvent);
            return ret;
        }

        if (!GetOverlappedResultEx(m_handle, &ovl, &nbytes, timeoutMs, TRUE)) {
            DWORD err = GetLastError();
            if (err == WAIT_TIMEOUT) {
                LOGE("Timed write to vioserial port: {} timeout", m_hostPortIdx);
                CancelIo(m_handle);
                ret = VioSOpStatus::TIMEOUT;
            } else if (err == ERROR_CANT_WAIT) {
                LOGE("Timed write to vioserial port: {} in disconnected state", m_hostPortIdx);
                CancelIo(m_handle);
                ret = VioSOpStatus::DISCONNECTED;
            } else {
                LOGE("Failed of timed write to vioserial port: {}, error: {}", m_hostPortIdx, err);
                ret = VioSOpStatus::FAILURE;
            }
        }
    }
    if (ret == VioSOpStatus::SUCCESS) {
        writeSz = nbytes;
    }
    CloseHandle(ovl.hEvent);

    return ret;
}

VioSOpStatus VioSAPI::WriteAllTimed(PVOID buf, DWORD writeSz, INT timeoutMs)
{
    DWORD nbytes;
    DWORD total = 0;
    VioSOpStatus ret;
    UINT8 *byteBuf = (UINT8 *)buf;

    do {
        nbytes = writeSz - total;
        ret = WriteTimed((PVOID)byteBuf, nbytes, timeoutMs);
        if (ret != VioSOpStatus::SUCCESS) {
            break;
        }
        total += nbytes;
        byteBuf += nbytes;
    } while (total < writeSz);
    return ret;
}

VioSOpStatus VioSAPI::ReportState()
{
    ViosConnHeader connHeader;
    uint32_t connHeaderSz = sizeof(connHeader);
    VioSOpStatus res;

    connHeader.magic = VIOS_PROTO_CONN_MAGIC;
    connHeader.len = 0;
    connHeader.opcode = VioSConnOpcode::REPORT_STATE;
    connHeader.state = m_state;
    res = WriteAllTimed(&connHeader, connHeaderSz, VIOS_CONN_TIMEOUT);
    if (res != VioSOpStatus::SUCCESS) {
        LOGE("Failed to send request to report state.");
        return res;
    }
    m_peerState = (VioSConnState)connHeader.state;

    return VioSOpStatus::SUCCESS;
}

VioSOpStatus VioSAPI::DropLeftData()
{
    uint8_t buff[256];
    VioSOpStatus res;
    while ((res = ReadAllTimed(buff, sizeof(buff), VIOS_CONN_TIMEOUT)) == VioSOpStatus::SUCCESS) {};
    return res;
}

VioSOpStatus VioSAPI::ReportGetPeerState()
{
    ViosConnHeader connHeader;
    uint32_t connHeaderSz = sizeof(connHeader);
    VioSOpStatus res;
    INT ntry = 0;

 send_retry:
    connHeader.magic = VIOS_PROTO_CONN_MAGIC;
    connHeader.len = 0;
    connHeader.opcode = VioSConnOpcode::GET_PEER_STATE;
    connHeader.state = m_state;
    res = WriteAllTimed(&connHeader, connHeaderSz, VIOS_CONN_TIMEOUT);
    if (res != VioSOpStatus::SUCCESS) {
        LOGE("Failed to send request to get peer state.");
        return res;
    }

    res = ReadAllTimed(&connHeader, connHeaderSz, VIOS_CONN_TIMEOUT);
    if (res != VioSOpStatus::SUCCESS) {
        LOGE("Cannot get state from peer.");
        m_peerState = VioSConnState::STATE_UNKNOWN;
        return res;
    }

    if ((connHeader.opcode != VioSConnOpcode::REPORT_STATE) ||
        (connHeader.state > VioSConnState::STATE_STOPPED)) {
        /* We may have some old data in read,
         * so drop it retry one more time
         */
        if (!ntry++) {
            DropLeftData();
            goto send_retry;
        }
        m_peerState = VioSConnState::STATE_BROKEN;
        return VioSOpStatus::BROKEN_CONN;
    }

    m_peerState = (VioSConnState)connHeader.state;
    return VioSOpStatus::SUCCESS;
}

VioSOpStatus VioSAPI::Disconnect()
{
    VioSOpStatus res;

    m_state = VioSConnState::STATE_STOPPED;
    res = ReportGetPeerState();
    if (res != VioSOpStatus::SUCCESS && m_peerState == VioSConnState::STATE_BROKEN) {
        m_state = VioSConnState::STATE_BROKEN;
        return VioSOpStatus::BROKEN_CONN;
    }

    return res;
}

VioSOpStatus VioSAPI::Connect(INT numRetry)
{
    VioSOpStatus res;
    DropLeftData();

    while (numRetry--) {
        m_state = VioSConnState::STATE_CONNECTING;
        m_peerState = VioSConnState::STATE_UNKNOWN;
        res = ReportGetPeerState();
        if (res != VioSOpStatus::SUCCESS) {
            if (res == VioSOpStatus::FAILURE) {
                m_state = VioSConnState::STATE_BROKEN;
                return VioSOpStatus::BROKEN_CONN;
            }
            if (m_peerState == VioSConnState::STATE_BROKEN) {
                m_state = VioSConnState::STATE_BROKEN;
                return VioSOpStatus::BROKEN_CONN;
            }
            Sleep(RECONNECT_RETRY_DELAY);
            continue;
        }

        if (m_state != VioSConnState::STATE_CONNECTING ||
            m_peerState != VioSConnState::STATE_CONNECTING) {
            return VioSOpStatus::FAILURE;
        }

        m_state = VioSConnState::STATE_CONNECTED;
        res = ReportGetPeerState();
        if (res != VioSOpStatus::SUCCESS) {
            if (m_peerState == VioSConnState::STATE_BROKEN) {
                m_state = VioSConnState::STATE_BROKEN;
                return VioSOpStatus::BROKEN_CONN;
            }
            continue;
        }
        if (m_state == VioSConnState::STATE_CONNECTED &&
            m_peerState == VioSConnState::STATE_CONNECTED) {
            return VioSOpStatus::SUCCESS;
        }
    }

    return VioSOpStatus::FAILURE;
}

VioSOpStatus VioSAPI::RecvPacket(PVOID pktBuf, DWORD pktBufSz, INT readTimeoutMs)
{
    ViosProtoHeader *packetHeader = (ViosProtoHeader*)pktBuf;
    uint32_t packetHeaderLen = sizeof(*packetHeader);
    uint8_t* packetData = (uint8_t*)packetHeader + packetHeaderLen;
    VioSOpStatus res;

    while (1) {
        if (m_state != VioSConnState::STATE_CONNECTED || m_peerState != VioSConnState::STATE_CONNECTED) {
            LOGE("Recv packet error: disconnected");
            return VioSOpStatus::FAILURE;
        }

        res = ReadAllTimed(packetHeader, packetHeaderLen, readTimeoutMs);
        if (res != VioSOpStatus::SUCCESS) {
            return res;
        }

        if (packetHeader->magic == VIOS_PROTO_CONN_MAGIC) {
            res = ProcessConnectPacket(packetHeader);
            if (res != VioSOpStatus::SUCCESS) {
                LOGE("Failed to connect");
                return res;
            }
            if (m_peerState == VioSConnState::STATE_CONNECTED) {
                continue;
            }
            return VioSOpStatus::FAILURE;
        }

        break;
    }

    if (packetHeader->len > (pktBufSz - packetHeaderLen)) {
        LOGE("Recv buffer left size {} too small to recv packet len {}",
             pktBufSz - packetHeaderLen, packetHeader->len);
        return VioSOpStatus::FAILURE;
    }

    res = ReadAllTimed(packetData, packetHeader->len, readTimeoutMs);
    if (res != VioSOpStatus::SUCCESS) {
        LOGE("Failed to receive packet data.");
        return res;
    }

    return VioSOpStatus::SUCCESS;
}

VioSOpStatus VioSAPI::ProcessConnectPacket(PVOID pktHeader)
{
    if (pktHeader == nullptr) {
        LOGE("Packet header is a null pointer.");
        return VioSOpStatus::FAILURE;
    }

    ViosConnHeader* connectPacketHeader = (ViosConnHeader*)pktHeader;
    if (connectPacketHeader->opcode != VIOS_CONN_OPCODE_GET_PEER_STATE) {
        LOGE("Peer state invalide opcode {}", connectPacketHeader->opcode);
        m_state = VioSConnState::STATE_BROKEN;
        m_peerState = VioSConnState::STATE_BROKEN;
        return VioSOpStatus::BROKEN_CONN;
    }

    m_peerState = (VioSConnState)connectPacketHeader->state;
    if (m_peerState == VioSConnState::STATE_BROKEN) {
        LOGE("Get broken state.");
        m_state = VioSConnState::STATE_BROKEN;
        return VioSOpStatus::BROKEN_CONN;
    }

    VioSOpStatus res = ReportState();
    if (res != VioSOpStatus::SUCCESS) {
        LOGE("Failed to report state.");
        m_state = VioSConnState::STATE_BROKEN;
        return VioSOpStatus::BROKEN_CONN;
    }

    return VioSOpStatus::SUCCESS;
}

VioSOpStatus VioSAPI::SendPacket(PVOID pktBuf, DWORD writeSz, INT totalTimeoutMs)
{
    if (m_state != VioSConnState::STATE_CONNECTED &&
        m_peerState != VioSConnState::STATE_CONNECTED) {
        LOGE("Failed to send packet, state {}, peer state {}", (INT)m_state, (INT)m_peerState);
        return VioSOpStatus::FAILURE;
    }

    return WriteAllTimed(pktBuf, writeSz, totalTimeoutMs);
}
